{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Broadcasts\n",
    "\n",
    "This notebook explains the different types of broadcast available in PyBaMM.\n",
    "Understanding of the [expression_tree](./expression-tree.ipynb) and [discretisation](../spatial_methods/finite-volumes.ipynb) notebooks is assumed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install \"pybamm[plot,cite]\" -q    # install PyBaMM if it is not installed\n",
    "import numpy as np\n",
    "\n",
    "import pybamm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also explicitly set up the discretisation that is used for this notebook. We use a small number of points in each domain, in order to easily visualise the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = pybamm.standard_spatial_vars\n",
    "geometry = {\n",
    "    \"negative electrode\": {var.x_n: {\"min\": pybamm.Scalar(0), \"max\": pybamm.Scalar(1)}},\n",
    "    \"negative particle\": {var.r_n: {\"min\": pybamm.Scalar(0), \"max\": pybamm.Scalar(1)}},\n",
    "}\n",
    "\n",
    "submesh_types = {\n",
    "    \"negative electrode\": pybamm.Uniform1DSubMesh,\n",
    "    \"negative particle\": pybamm.Uniform1DSubMesh,\n",
    "}\n",
    "\n",
    "var_pts = {var.x_n: 5, var.r_n: 3}\n",
    "mesh = pybamm.Mesh(geometry, submesh_types, var_pts)\n",
    "\n",
    "spatial_methods = {\n",
    "    \"negative electrode\": pybamm.FiniteVolume(),\n",
    "    \"negative particle\": pybamm.FiniteVolume(),\n",
    "}\n",
    "disc = pybamm.Discretisation(mesh, spatial_methods)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Primary broadcasts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Primary broadcasts are used to broadcast from a \"larger\" scale to a \"smaller\" scale, for example broadcasting temperature T(x) from the electrode to the particles, or broadcasting current collector current i(y, z) from the current collector to the electrodes.\n",
    "To demonstrate this, we first create a variable `T` on the negative electrode domain, discretise it, and evaluate it with a simple linear vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  ],\n",
       "       [0.25],\n",
       "       [0.5 ],\n",
       "       [0.75],\n",
       "       [1.  ]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "T = pybamm.Variable(\"T\", domain=\"negative electrode\")\n",
    "disc.set_variable_slices([T])\n",
    "disc_T = disc.process_symbol(T)\n",
    "disc_T.evaluate(y=np.linspace(0, 1, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then broadcast `T` onto the \"negative particle\" domain (using primary broadcast as we are going from the larger electrode scale to the smaller particle scale), and discretise and evaluate the resulting object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.  ],\n",
       "       [0.  ],\n",
       "       [0.  ],\n",
       "       [0.25],\n",
       "       [0.25],\n",
       "       [0.25],\n",
       "       [0.5 ],\n",
       "       [0.5 ],\n",
       "       [0.5 ],\n",
       "       [0.75],\n",
       "       [0.75],\n",
       "       [0.75],\n",
       "       [1.  ],\n",
       "       [1.  ],\n",
       "       [1.  ]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "primary_broad_T = pybamm.PrimaryBroadcast(T, \"negative particle\")\n",
    "disc_T = disc.process_symbol(primary_broad_T)\n",
    "disc_T.evaluate(y=np.linspace(0, 1, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The broadcasted object makes 3 (since the r-grid has 3 points) copies of each element of `T` and stacks them all up to give an object with size 3x5=15. In the resulting vector, the first 3 entries correspond to the 3 points in the r-domain at the first x-grid point (where T=0 uniformly in r), the next 3 entries correspond to the next 3 points in the r-domain at the second x-grid point (where T=0.25 uniformly in r), etc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Secondary broadcasts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Secondary broadcasts are used to broadcast from a \"smaller\" scale to a \"larger\" scale, for example broadcasting SPM particle concentrations c_s(r) from the particles to the electrodes. Note that this wouldn't be used to broadcast particle concentrations in the DFN, since these already depend on both x and r.\n",
    "To demonstrate this, we first create a variable `c_s` on the negative particle domain, discretise it, and evaluate it with a simple linear vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0. ],\n",
       "       [0.5],\n",
       "       [1. ]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c_s = pybamm.Variable(\"c_s\", domain=\"negative particle\")\n",
    "disc.set_variable_slices([c_s])\n",
    "disc_c_s = disc.process_symbol(c_s)\n",
    "disc_c_s.evaluate(y=np.linspace(0, 1, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then broadcast `c_s` onto the \"negative electrode\" domain (using secondary broadcast as we are going from the smaller particle scale to the large electrode scale), and discretise and evaluate the resulting object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0. ],\n",
       "       [0.5],\n",
       "       [1. ],\n",
       "       [0. ],\n",
       "       [0.5],\n",
       "       [1. ],\n",
       "       [0. ],\n",
       "       [0.5],\n",
       "       [1. ],\n",
       "       [0. ],\n",
       "       [0.5],\n",
       "       [1. ],\n",
       "       [0. ],\n",
       "       [0.5],\n",
       "       [1. ]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "secondary_broad_c_s = pybamm.SecondaryBroadcast(c_s, \"negative electrode\")\n",
    "disc_broad_c_s = disc.process_symbol(secondary_broad_c_s)\n",
    "disc_broad_c_s.evaluate(y=np.linspace(0, 1, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The broadcasted object makes 5 (since the x-grid has 5 points) identical copies of the whole variable `c_s` to give an object with size 5x3=15. In the resulting vector, the first 3 entries correspond to the 3 points in the r-domain at the first x-grid point (where c_s varies in r), the next 3 entries correspond to the next 3 points in the r-domain at the second x-grid point (where c_s varies in r), etc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "The relevant papers for this notebook are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, and others. Array programming with NumPy. Nature, 585(7825):357–362, 2020. doi:10.1038/s41586-020-2649-2.\n",
      "[2] Valentin Sulzer, Scott G. Marquis, Robert Timms, Martin Robinson, and S. Jon Chapman. Python Battery Mathematical Modelling (PyBaMM). Journal of Open Research Software, 9(1):14, 2021. doi:10.5334/jors.309.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pybamm.print_citations()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
